Adding metal kernels for the gated delta nets. - #4020

Hi @tpegolotti — thanks for this PR, and thanks @zcbenz for pointing me at it. The chunked WY-form kernels are a clear step up for batched prefill (the B ≥ 4 curves speak for themselves), and I'd be glad to see this land. @zcbenz asked me to look for possible improvements, so here are my notes. Everything below comes from reading the diff — I have not run the branch yet — so please take it as suggestions and correct me wherever I've misread something. I'm happy to help with any of it (patches against this branch, or measurements on an M5 Max).

1. gated_delta_seq layout (which is also the baseline of the speedup plots). This is the one-SIMD-group-per-value-row kernel from mlx-lm: at Dk = 128 each lane holds 4 state elements and every timestep does two full 32-lane simd_sums. In Add packed gated delta kernel, bitwise-pinned by an explicit-tree comparator mlx-lm#1559 I pack 8 value rows per SIMD-group (4 lanes per row, 32 contiguous state elements per lane, simd_shuffle_xor(1, 2) for the last two reduction levels); it is bitwise-identical to the current kernel and ~1.9–2.0× faster at B = 1, T = 2048–8192 on M5 (independently reproduced in that PR). Reading off your M5 Max plot, B = 1 sits at roughly 1.3–1.45× (simdgroup) and 1.6–2.2× (NAX) for T ≤ 2048, so the packed layout seems to be in the same range at B = 1 — that's a different harness though, so it needs a proper head-to-head before anyone relies on it, and I'd be glad to run that. Independently of the B = 1 question, the port is mechanical and would make the T ≤ 16 path (decode / short prefill) faster while keeping it bitwise-equal to today's kernel. I'd be happy to send it as a PR against your branch if you'd like.

2. Gate precision. gates / beta are cast to out_dtype before the primitive, so for bf16 models the decay is quantized to bf16 before the log + prefix sum. mlx-lm computes g in float32 today (compute_g), so this would be a small precision regression relative to the current mlx-lm path. Since g / beta are only [B, T, H], keeping them in fp32 inside the op looks cheap.

3. dtype guard. Only float and bfloat16_t are instantiated, and use_fallback doesn't look at the dtype, so I believe float16 inputs would look up seq_gated_delta_float16_t_… and fail at kernel lookup. Either instantiating float16_t or routing fp16 to the fallback would cover it.

4. Vector-gate guard. use_fallback doesn't check g.ndim. The fallback handles a [B, T, H, Dk] gate, but when the shape checks pass, such a gate reaches the kernel, which indexes g as [B, T, Hv] — so the result would be silently wrong rather than raising. (mlx-lm's kimi_linear does pass a vector gate; today it only stays out of the kernel because its Dk is 72.)

5. Head list. supported_heads = {(24,24), (32,32), (16,32), (16,48)} misses Qwen3.5-0.8B / 2B (Hk, Hv = 16, 16) and Qwen3.5-122B-A10B / 397B-A17B (16, 64), which would then take the per-timestep fallback (roughly 20–25 graph ops per timestep, i.e. tens of thousands of nodes per layer at T = 2048 — much slower than the current mlx-lm kernel). As far as I can see Hk / Hv only feed strides and the GQA index, so they could be runtime arguments; only Dk / Dv need to be compile-time.

6. Masks take the same fallback. mlx-lm's batched generation passes padding masks, so I think the masked path needs a kernel before mlx-lm can switch over wholesale (the existing mlx-lm masked kernel is a ready reference, and I'm happy to help there too).

7. Smaller things:

   • fill_gpu zero-fills out on every call, but all three kernels write every element, so that's an extra pass + dispatch per layer per token.
   • std::getenv / std::stoi run inside eval_gpu on the decode hot path; caching them in a static would avoid that.
   • The gamma[] threadgroup hand-off (lanes < C write, all lanes read) has no simdgroup_barrier; a simd_shuffle of gamma_val would avoid threadgroup memory altogether.
   • metal::fast::log / exp plus the max(g, 1e-6) clamp change the numerics relative to the sequential kernel, and the NAX test needs atol=1e-1 where the others pass at 1e-4. It might be worth a token-identity check on a real model (greedy decode, sequential vs chunked) in addition to the wikitext perplexity.
   • Docs: the Python docstring shapes are transposed relative to the implementation ([B, H, T, Dk] / [B, H, Dk, Dv] vs actual [B, T, H, Dk] / [B, Hv, Dv, Dk]), gamma is undocumented (it's the linear-space decay, not its log), state() is still a TODO, LOAD_M / LOAD_MT are defined twice in gated_delta_update_impl.h, and there's a duplicated return false; in use_fallback.

Again, happy to turn any of these into patches or measurements — just say which would be most useful. (Also glad to help with the rebase onto current main if that's useful.)

Hi @wyanzhao, thanks for the comments and suggestions! I've added most of them in the latest commits.

1. Layout: Yes, please send your PR.
2. Gate precision: Removing the gamma / beta casts broke the PPL evaluation.
3. Dtype: Added the float16_t instantiation
4. Gate guard: Inside the fallback, I expand the dimensions with a check on the number of dimensions. Is this what you were looking for?
5. Head list: Added the dimensions for the additional Qwen models.
6. Masks: Masks are currently supported only in the fallback. It's in the TODOs for the kernels.
7. Smaller things:
   • Remove the fill_gpu in the eval_gpu call
   • Reading the environment variable is still needed to swap kernels dynamically for testing purposes. I will remove it in the future.
   • Added a simdgroup barrier after initializing gamma.
   • The NAX test has lower tolerance due to the relaxed precision of the matrix multiplication, whereas all other kernels compute the intermediate results in float32.
   • Updated the docs and removed the duplicate code.

Thanks for the invitation — the packed-seq PR is up at #4409.

Short version: bitwise-equal by construction; sequential is 1.86–2.16×
at long T on M5 Max, and T ≤ 16 is a small lift with no regression
beyond noise.

One heads-up: this makes the sequential denominator of your speedup
plots faster, so those numbers will move and will need a recompute or
a note.

I left two pre-existing items out of that PR: use_fallback still
doesn't admit the newly instantiated (16, 16) / (16, 64) pairs, and
there's still no g.ndim guard (a 4-D vector gate that passes the
shape checks would be read as [B, T, Hv]). Those are out of scope
here; if you want them, a small follow-up patch from you would cover
them.